#!/usr/bin/env python3
"""
统计卡片字段迁移脚本
为公司信息表添加统计卡片相关字段
"""

import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from sqlalchemy import text
from app.database import engine, SessionLocal

def migrate_stat_cards():
    """迁移公司信息表，添加统计卡片字段"""
    
    # 要添加的新字段
    new_columns = [
        "stat1_title VARCHAR(100) DEFAULT '年行业经验'",
        "stat1_value VARCHAR(50)",
        "stat2_title VARCHAR(100) DEFAULT '完成项目'", 
        "stat2_value VARCHAR(50)",
        "stat3_title VARCHAR(100) DEFAULT '专业团队'",
        "stat3_value VARCHAR(50)",
        "stat4_title VARCHAR(100) DEFAULT '客户满意度'",
        "stat4_value VARCHAR(50) DEFAULT '100%'"
    ]
    
    db = SessionLocal()
    
    try:
        # 检查表是否存在
        result = db.execute(text("SELECT name FROM sqlite_master WHERE type='table' AND name='company_info'"))
        if not result.fetchone():
            print("company_info 表不存在，请先运行数据库初始化")
            return
        
        # 检查字段是否已存在
        for column_def in new_columns:
            column_name = column_def.split()[0]
            try:
                result = db.execute(text(f"PRAGMA table_info(company_info)"))
                columns = [row[1] for row in result.fetchall()]
                
                if column_name not in columns:
                    # 添加字段
                    db.execute(text(f"ALTER TABLE company_info ADD COLUMN {column_def}"))
                    print(f"添加字段: {column_name}")
                else:
                    print(f"字段 {column_name} 已存在，跳过")
            except Exception as e:
                print(f"添加字段 {column_name} 失败: {e}")
        
        # 为现有记录设置默认值
        try:
            # 检查是否有现有记录
            result = db.execute(text("SELECT COUNT(*) FROM company_info"))
            count = result.scalar()
            
            if count > 0:
                # 更新现有记录的默认值
                update_sql = """
                UPDATE company_info SET 
                    stat1_title = COALESCE(stat1_title, '年行业经验'),
                    stat1_value = COALESCE(stat1_value, '25+'),
                    stat2_title = COALESCE(stat2_title, '完成项目'),
                    stat2_value = COALESCE(stat2_value, '500'),
                    stat3_title = COALESCE(stat3_title, '专业团队'),
                    stat3_value = COALESCE(stat3_value, '50-100人'),
                    stat4_title = COALESCE(stat4_title, '客户满意度'),
                    stat4_value = COALESCE(stat4_value, '100%')
                WHERE stat1_title IS NULL OR stat1_value IS NULL 
                   OR stat2_title IS NULL OR stat2_value IS NULL
                   OR stat3_title IS NULL OR stat3_value IS NULL
                   OR stat4_title IS NULL OR stat4_value IS NULL
                """
                db.execute(text(update_sql))
                print("更新现有记录的统计卡片默认值")
        except Exception as e:
            print(f"更新默认值失败: {e}")
        
        db.commit()
        print("统计卡片字段迁移完成！")
        
    except Exception as e:
        print(f"迁移失败: {e}")
        db.rollback()
    finally:
        db.close()

if __name__ == "__main__":
    migrate_stat_cards()
